from flask import Flask, request, jsonify
import torch
import numpy as np
from config import Config
from model.dnn import DNNModel

app = Flask(__name__)
config = Config()

# 加载模型和索引
model = DNNModel(config)
model.load_state_dict(torch.load(config.DNN_MODEL_SAVE_PATH))
model.eval()

"""
response = requests.get(
    "http://localhost:5000/recommend",
    params={"user_id": 123, "item_id": 456}
)
print(response.json())
"""
@app.route('/recommend', methods=['GET'])
def recommend():
    try:
        # 获取请求参数
        user_id = int(request.args.get('user_id'))
        item_id = int(request.args.get('item_id'))
        
        # 推理
        with torch.no_grad():
            scores = model(torch.tensor([user_id]),torch.tensor([item_id]))
            #scores = model(user_id,item_id) #错误，model神经网络的输入必须为张量
        

        # 构建响应
        return jsonify({
            "user_id": user_id,
            "item_id": item_id,
            "score": scores.item()
        })
    
    except Exception as e:
        return jsonify({"error": str(e)}), 400
"""
batch_response = requests.post(
    "http://localhost:5000/batch_recommend",
    json={ "user_id": 2344, "item_ids": [123,456,789,3423,234],}
)
print(batch_response.json())
"""

@app.route('/batch_recommend', methods=['POST'])
def batch_recommend():
    try:
        # 获取批量用户ID
        data = request.json
        user_id = data['user_id']
        item_ids = data['item_ids']
        
        # 批量推理        
        results = []
        for i, item_id in enumerate(item_ids):
            with torch.no_grad():
                scores = model(torch.tensor([user_id]),torch.tensor([item_id]))
            results.append({
                "item_id": item_id,
                "score": float(scores.item())
            })
        
        return jsonify({"user_id": user_id,"results": results})
    
    except Exception as e:
        return jsonify({"error": str(e)}), 400

if __name__ == '__main__':
    app.run(host=config.API_HOST, port=config.API_PORT, threaded=True)
